import json
from .tushare_api import TushareAPI
from .akshare_api import AkshareAPI
from .tushare_api_2 import TushareAPI_stat
from .tushare_api_3 import TushareAPI_Emotion
from .tushare_api_4 import TushareAPI_Reference
from channels.generic.websocket import AsyncWebsocketConsumer
from data_analysis.tasks import start_analysis_tasks
'''
处理Tushare数据更新，并返回数据更新信息以及进度
'''
class TushareProgressConsumer(AsyncWebsocketConsumer):
    """
    Tushare数据更新进度WebSocket消费者
    继承自通用进度更新消费者
    """

    async def connect(self):
        print(f'try to connect websocket!')
        """建立WebSocket连接"""
        await self.accept()

    async def disconnect(self, close_code):
        """断开WebSocket连接"""
        pass

    async def send_progress(self, current, total, message):
        """
        发送进度更新消息
        
        Args:
            current: 当前进度
            total: 总数
            message: 进度消息
        """
        await self.send(text_data=json.dumps({
            'type': 'progress',
            'current': current,
            'total': total,
            'message': message
        }))
    
    async def send_complete(self, message="任务完成"):
        """
        发送任务完成消息
        
        Args:
            message: 完成消息
        """
        await self.send(text_data=json.dumps({
            'type': 'complete',
            'message': message
        }))
    
    async def send_error(self, message):
        """
        发送错误消息
        
        Args:
            message: 错误消息
        """
        await self.send(text_data=json.dumps({
            'type': 'error',
            'message': message
        }))
    async def receive(self, text_data):
        """
        接收WebSocket消息并处理数据更新请求
        """
        try:
            data = json.loads(text_data)
            if data['type'] == 'start_update':
                # 创建API实例
                tushare_api = TushareAPI()
                tushare_api_stat = TushareAPI_stat()
                tushare_api_emotion = TushareAPI_Emotion()
                tushare_api_ref = TushareAPI_Reference()
                akshare_api = AkshareAPI()

                
                try:
                    # 发送开始消息
                    await self.send_progress(
                        current=1,
                        total=5,
                        message='开始更新Tushare基础数据和行情数据...'
                    )

                    # 先更新Tushare数据的基础数据和行情数据
                    await tushare_api.update_all_data(
                        start_date=data['start_date'],
                        end_date=data['end_date'],
                        websocket=self
                    )
                    # 发送开始消息
                    await self.send_progress(
                        current=2,
                        total=5,
                        message='开始更新Tushare统计类数据...'
                    )
                    # 再更新Tushare数据的统计数据
                    await tushare_api_stat.update_all_data(
                        start_date=data['start_date'],
                        end_date=data['end_date'],
                        websocket=self
                    )
                    # 发送开始消息
                    await self.send_progress(
                        current=3,
                        total=5,
                        message='开始更新Tushare打板专题数据...'
                    )
                    # 再更新Tushare数据的打板专题数据
                    await tushare_api_emotion.update_all_data(
                        start_date=data['start_date'],
                        end_date=data['end_date'],
                        websocket=self
                    )
                    # 发送开始消息
                    await self.send_progress(
                        current=4,
                        total=5,
                        message='开始更新Tushare参考数据...'
                    )
                    await tushare_api_ref.update_all_data(
                        start_date=data['start_date'],
                        end_date=data['end_date'],
                        websocket=self
                    )

                    # 更新进度
                    await self.send_progress(
                        current=5,
                        total=5,
                        message='Tushare数据更新完成，开始更新Akshare数据...'
                    )

                    # 再更新Akshare数据
                    await akshare_api.update_all_data(
                        start_date=data['start_date'],
                        end_date=data['end_date'],
                        websocket=self
                    )

                    # 执行数据分析
                    try:
                        # analyze_market_data(period='day')
                        start_analysis_tasks()
                        print(f"市场数据分析已完成")
                    except Exception as e:
                        await self.send_error(f'数据分析失败: {str(e)}')
                        return

                    # 发送完成消息
                    await self.send_complete('所有数据更新和分析完成')

                except Exception as e:
                    # 发送错误消息
                    await self.send_error(f'数据更新失败: {str(e)}')

        except Exception as e:
            # 发送错误消息
            await self.send_error(f'消息处理失败: {str(e)}')